package com.sui.bigdata.sml.web.controller;

import com.alibaba.fastjson.JSON;
import com.sui.bigdata.sml.web.dto.ConvertModelDTO;
import com.sui.bigdata.sml.web.dto.ModelDTO;
import com.sui.bigdata.sml.web.dto.ModelGroupDTO;
import com.sui.bigdata.sml.web.service.ModelGroupService;
import com.sui.bigdata.sml.web.service.ModelService;
import com.sui.bigdata.sml.web.util.Constants;
import com.sui.bigdata.sml.web.util.ModelType;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.bind.annotation.*;

import javax.validation.Valid;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * @author songhaicheng
 * @date 2019/11/6 10:50
 * @description
 * @reviewer
 */
@RestController
@CrossOrigin
@RequestMapping(value = "/web/ui/models")
@Slf4j
public class ModelController {
    private static final Logger LOGGER = LoggerFactory.getLogger(ModelController.class);
    @Autowired
    private ModelService service;
    @Autowired
    private ModelGroupService groupService;

    @PostMapping("/convert/spark-to-pmml")
    ResponseEntity convertSparkToPmml(@Valid @RequestBody ConvertModelDTO convert) {
        convert.setType(ModelType.PMML.getValue());
        if (service.convertSparkModels(convert)) {
            return new ResponseEntity<>(convert.getDestination(), HttpStatus.OK);
        } else {
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    @PostMapping("/convert/spark-to-mleap")
    ResponseEntity convertSparkToMleap(@Valid @RequestBody ConvertModelDTO convert) {
        LOGGER.info("Spark2Mleap接口请求参数为:{}", JSON.toJSONString(convert));
        convert.setType(ModelType.MLEAP.getValue());
        if (service.convertSparkModels(convert)) {
            return new ResponseEntity<>(convert.getDestination(), HttpStatus.OK);
        } else {
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    @GetMapping
    ResponseEntity listByUserId() {
        return new ResponseEntity<>(service.listByUserId(Constants.DEFAULT_USER_ID), HttpStatus.OK);
    }

    @PostMapping
    ResponseEntity add(@Valid @RequestBody ModelDTO dto) {
        if (dto.getFields() == null || dto.getFields().size() == 0) {
            throw new RuntimeException("字段信息不能为空");
        }
        // 所有新增的模型 version 都是 1
        dto.setVersion(1);
        dto.setId(service.add(dto));
        log.info("模型新增成功：{}", dto.toString());
        return new ResponseEntity<>(dto, HttpStatus.CREATED);
    }

    @GetMapping("/{token}")
    ResponseEntity listByToken(@PathVariable("token") String token) {
        return new ResponseEntity<>(service.listByToken(token), HttpStatus.OK);
    }

    @Transactional
    @PutMapping("/{token}")
    ResponseEntity push(@PathVariable("token") String token, @Valid @RequestBody ModelDTO dto) {
        dto.setToken(token);
        // 更新模型，把 version 置空，交给 service 进行下一步处理
        dto.setVersion(null);
        dto.setId(service.add(dto));
        log.info("模型更新成功：{}", dto.toString());
        refreshServingByModel(dto.getToken());
        return new ResponseEntity<>(dto, HttpStatus.OK);
    }

    @Transactional
    @PutMapping("/{token}/info")
    ResponseEntity updateModelDescription(@PathVariable("token") String token, @RequestBody ModelDTO dto) {
        ModelDTO modelDTO = service.get(dto.getId());
        modelDTO.setDescription(dto.getDescription());
        if (service.update(modelDTO)) {
            return new ResponseEntity<>(dto, HttpStatus.OK);
        } else {
            log.error("更新模型描述信息失败！");
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    @DeleteMapping("/{token}")
    ResponseEntity delete(@PathVariable("token") String token) {
        List<ModelGroupDTO> havingGroups = groupService.listHavingModelToken(token);
        if (havingGroups.size() > 0) {
            throw new RuntimeException("检测到有 " + havingGroups.size() + " 个线上服务（" +
                    havingGroups.stream().map(ModelGroupDTO::getToken).collect(Collectors.joining(", "))
                    + "）正在使用该模型，请先下线相关服务");
        }
        if (service.delete(token)) {
            return new ResponseEntity<>(null, HttpStatus.OK);
        } else {
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    @Transactional
    @PutMapping("/{token}/switch/{version}")
    ResponseEntity switchVersion(@PathVariable("token") String token, @PathVariable("version") Integer version) {
        if (service.switchVersion(token, version)) {
            log.info("模型版本指定成功，token：{}，version：{}", token, version);
            refreshServingByModel(token);
            return new ResponseEntity<>(null, HttpStatus.OK);
        } else {
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    @PostMapping("/serving")
    ResponseEntity addServing(@Valid @RequestBody ModelGroupDTO groupDTO) {
        ModelGroupDTO dto = groupService.add(groupDTO);
        log.info("保存模型组成功：{}", dto.toString());
        return new ResponseEntity<>(dto, HttpStatus.OK);
    }

    @GetMapping("/serving")
    ResponseEntity listServingByUserId() {
        return new ResponseEntity<>(groupService.listByUserId(Constants.DEFAULT_USER_ID), HttpStatus.OK);
    }

    @GetMapping("/serving/{token}")
    ResponseEntity getServingInfo(@PathVariable("token") String servingToken) {
        return new ResponseEntity<>(groupService.getInfoByServingToken(servingToken), HttpStatus.OK);
    }

    @Transactional
    @PutMapping("/serving/{token}")
    ResponseEntity updateGroup(@PathVariable("token") String servingToken, @RequestBody List<String> models) {
        // 此接口不接受对 group 激活状态的修改
        if (groupService.modifyServingModels(servingToken, models)) {
            log.info("更新模型组成功：{}", models);
            if (groupService.getByServingToken(servingToken).getActive()) {
                groupService.refreshServing(servingToken);
            }
            return new ResponseEntity<>(null, HttpStatus.OK);
        } else {
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    @PostMapping("/serving/{token}/offline")
    ResponseEntity servingOff(@PathVariable("token") String servingToken) {
        if (groupService.modifyServingActive(servingToken, false)) {
            log.warn("下线模型组成功：{}", servingToken);
            return new ResponseEntity<>(null, HttpStatus.OK);
        } else {
            log.error("下线模型组失败：{}", servingToken);
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    @PostMapping("/serving/{token}/online")
    ResponseEntity servingOn(@PathVariable("token") String servingToken) {
        if (groupService.modifyServingActive(servingToken, true)) {
            log.warn("激活模型组成功：{}", servingToken);
            return new ResponseEntity<>(null, HttpStatus.OK);
        } else {
            log.error("激活模型组失败：{}", servingToken);
            return new ResponseEntity<>(null, HttpStatus.INTERNAL_SERVER_ERROR);
        }
    }

    private void refreshServingByModel(String modelToken) {
        // 更新线上模型服务组
        AtomicInteger refreshed = new AtomicInteger();
        groupService.listHavingModelToken(modelToken).parallelStream().forEach(g -> {
            if (groupService.refreshServing(g.getToken())) {
                refreshed.getAndIncrement();
            }
        });
        log.info("检测到模型 {} 版本有更新，成功刷新用到此模型的线上服务 {} 个", modelToken, refreshed.get());
    }
}
